import cv2
from pathlib import Path
from shutil import copy


_video_types = {".avi", ".mp4", ".mkv"}


def _vaild_src_and_dst(src: Path, dst: Path):
    video_types = _video_types
    if not (src.is_file() and src.suffix in video_types) or not (dst.is_dir()):
        raise ValueError("输入文件不是视频，或者输出目录不是文件夹")


def _get_frame_and_print(src: str, need_print=None):
    cap = cv2.VideoCapture(src)
    total_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT)
    if total_frames == 0:
        raise ValueError("输入文件不是视频")
    processed_frame = 0

    print(f"total: {total_frames} frames processing...")
    while True:
        succ, frame = cap.read()
        if not succ:
            print(f"视频拆帧完毕, 一共拆了{processed_frame}帧")
            return

        yield frame, processed_frame

        if need_print:
            print(f"{need_print}{int(round(processed_frame/total_frames, 2) * 100)}")
        processed_frame += 1


def do_split(src: Path, dst: Path, need_print=None):
    _vaild_src_and_dst(src, dst)
    frame_reader = _get_frame_and_print(str(src), need_print)
    for frame, f_count in frame_reader:
        cv2.imencode(".png", frame)[1].tofile(dst / f"{f_count:08}.png")


def do_split_with_skip_frame(
    src: Path, dst: Path, gt: Path, skip_frames: int, need_print=None
):
    choosed_gt = gt / "selected"
    choosed_gt.mkdir(exist_ok=True)
    skip_target = skip_frames + 1 if skip_frames >= 0 else 2
    frame_reader = _get_frame_and_print(str(src), need_print)
    for frame, count in frame_reader:
        if count % skip_target:
            continue

        cv2.imencode(".png", frame)[1].tofile(dst / f"{count:08}.png")
        txt_file_name = f"{count:08}.txt"
        copy(gt / txt_file_name, choosed_gt / txt_file_name)
